from XSY.DEVS.CoupledDEVS import *
from XSY.DEVS.FSDEVS import *
from XSY.hashdict import *
from copy import deepcopy
#from XSY.GraphVizy.GraphVizy import *
from XSY.Graph.Graph import *
from XSY.Set.TID import * 
#from odict import odict
from datetime import datetime

def printV(s, V):
   '''printing utility function'''
   print s, len(V)
   for v in V:
      print str(v)

def get_schedules(discstate):
   '''get schedules from a discreate state'''
   schedules = hashdict()#odict()
   for key in discstate:
      ta = discstate[key][1]#['ts']
      if isnumeric(ta):
         schedules[key] = Interval(ta, ta)
      elif isinstance(ta, Interval):
         schedules[key] = ta
      elif isinstance(ta, tuple):
         schedules[key] = Interval(ta[0],ta[1])

   return schedules

def isSameStateVectors(svA, svB):
   '''check if state vector svA is equal to svB'''
   for k in svA:
      if k in svB:
         if not( svA[k][0] == svB[k][0] and svA[k][1] == svB[k][1]):
            return False   
      else:
         return False
   return True

class ReachableVtx:
   '''if v is an instance of ReachableVtx, then v.s is the discrete state that is a dictionary, v.tz is the time zone'''
   ss_value_repo = {} # state schedule repository for efficient memory management
   ss_key_repo = {}   # repository of keys for state_schedules
   vtx_s_repo = {}
   vtx_tz_repo = {}

   def __init__(self, s, tz):
      '''s is a discrete state, tz ia a time zone'''
      if isinstance(s, dict):
         self.s = hashdict()
         for k in s:   # discrete state dictionary
            self.s[k]=deepcopy(s[k])
            #self.set_ss(k, s[k])
         #print "type=",type(self.s)
         #print "self.s=%s"%self.s
         #self.s.sort()
         #print "self.s=%s"%self.s

         #raw_input("ok?")
      else:
         raise Exception("%s is not a dictionary but %s!"%(s,classname(s)))

      if isinstance(tz, TID):
         self.tz = TID().copy(tz) # time zone of TID
      else:
         raise Exception("%s is not a TID but %s!"%(tz,classname(tz)))

 
   def set_ss(self, hn, value):
      ''' set state schedule [hn] = value using their corresponding stored cache'''
      if hn not in ReachableVtx.ss_key_repo:
         ReachableVtx.ss_key_repo[hn] = hn
      hnkey = ReachableVtx.ss_key_repo[hn]

      if value not in ReachableVtx.ss_value_repo:
         ReachableVtx.ss_value_repo[value] = value
      ssvalue = ReachableVtx.ss_value_repo[value]
      self.s[hnkey] = ssvalue

   def use_repo(self):
      '''use memory repository for ReachableVtx'''
      for k in self.s:
         self.set_ss(k, self.s[k])

      self.tz.use_repo()

      if self.s not in self.vtx_s_repo:
         self.vtx_s_repo[self.s] = self.s
      ss = self.vtx_s_repo[self.s]

      if self.tz not in self.vtx_tz_repo:
         self.vtx_tz_repo[self.tz] = self.tz
      tz = self.vtx_tz_repo[self.tz]

      self.s = ss
      self.tz = tz

   def __str__(self):
      return str(self.s)+","+str(self.tz)

   def __hash__(self):
      h = 0
      for k in self.s:
         h += hash(k)*hash(self.s[k][0])#['s'])
         h += hash(k)*hash(self.s[k][1])#['ts'])
         h += hash(k)*hash(self.tz)
      return h
    
   def __eq__(self, ob):
      return str(self)==str(ob) 

   @staticmethod
   def PrintRepository():
      print 'ReachableVtx: |ss_value_repo|=%d, |ss_key_repo|=%d, |vtx_ss_repo|=%d, |vtx_tz_repo|=%d'%\
            (len(ReachableVtx.ss_value_repo),len(ReachableVtx.ss_key_repo),  len(ReachableVtx.vtx_s_repo), len(ReachableVtx.vtx_tz_repo))

def copysetto(A, B):
   '''copy a set A to a set B'''
   if isinstance(B, list):
      for a in A:
         B.append(a)
   elif isinstance(B, set):
      for a in A:
         B.add(a)


class Reachability:
   '''reachability graph generator'''
   edge_CR_repo = {}
   edge_trg_z_repo = {}
   hiername_repo = {}
   edge_p_repo={}
   sTrg='trg'
   sp='p'
   sr='r'

   def PrintRepository(self):
      print '%s: |edge_CR_repo|=%d, |edge_trg_z_repo|=%d, |hiername_repo|=%d, |edge_p_repo|=%d'% \
       (classname(self), len(self.edge_CR_repo), len(self.edge_trg_z_repo), len(self.hiername_repo), len(self.edge_p_repo))
      ReachableVtx.PrintRepository()
      TID.PrintRepository()

   def __init__(self, md, X):
      '''constructor'''
      self.AM={}
 
      if isinstance(md, CoupledDEVS) or isinstance(md, FSDEVS):
         self.G = self.GenReachFromFSD(md,X)
      else:
         raise Exception("Invalide input model:%s!"%md)
      print '\nReachability Graph: |V|=%d, |E|=%d'%(len(self.G.V), self.G.lenOfE())

   def getDRV(self, rv):
      ''' '''
      if isinstance(rv, DRV):
         return rv
      else:
         return DRV({rv:1})

   def get_ta(self, state, hn):
      '''get ta'''
      if isinstance(hn, str):
         md = self.AM[hn]
      elif isinstance(hn, FSDEVS):
         md = hn

      md.set_state(state)
      rv_ta =md.rv_tau()
      if isinstance(rv_ta, TInterval):
         return (rv_ta.l, rv_ta.u)
      else:
         rv_ta = self.getDRV(rv_ta)
         if len(rv_ta.table)==1:
            ta = rv_ta.sample()
            return ta
         else:
            raise Exception("Invalid rv_ta in %s"%function_name())

   def Delta0(self):
      '''get a list of initial state vectors'''
      discstates=hashdict() #odict()

      for hn in self.AM: # for each atomic model
         rv_s0 = self.getDRV(self.AM[hn].rv_delta_0()).table
         discs=[]
         for s0 in rv_s0:
            rv_ta = self.get_ta(s0, hn) 
            discs.append((s0, rv_ta))#hashdict({'s':s0, 'ts':rv_ta}))

         discstates[hn]=[]
         copysetto(discs, discstates[hn])

      return cross_product_dict(discstates, hashdict())#odict()) #hashdict())

   def set_atomic_model_state(self, state_vector):
      '''set the state of md by s. If s doesn't contains n, s=v, otherwise s=v,n'''
      for hn in state_vector:
         md = self.AM[hn]
         #print 'set_state[%s]=%s'%(hn, state_vector[hn]['s'])
         md.set_state(state_vector[hn][0])#['s'])

   def IsEnabledY(self, tz, ts, hn):
      '''check if hn is enabled to generate output and internal transition'''
      if tz.valid() == False:
         return False
      else:
         if isnumeric(ts): ta = Interval(ts, ts)
         elif isinstance(ts, Interval): ta = ts
         elif isinstance(ts, tuple): ta = Interval(ts[0], ts[1])
         if ta.infinity(): return False
         else: 
            inter = tz.b(hn).intersect(ta)
            return inter.valid()

#   def Z_yx(y):
#      ''''''
#      XP = set()
#      if isinstance(y, PortValue):
#         md = y.port.parent
#         if isinstance(md.parent, CoupledDEVS):
#            md.parent.getAllInfluencee(y, XP)
#      return XP
   
   def hiername(self, md):
      '''get hierarchical name of md'''
      hn = hiername(md)
      if hn not in self.hiername_repo:
         self.hiername_repo[hn]=hn
      answer = self.hiername_repo[hn]
      return answer

   def get_state_vector_of_influencee(self, discs, z):
      crossprd = []
      if z and z.port.parent:
         md = z.port.parent
         hn = self.hiername(md)
         self.set_atomic_model_state(discs)  
         crossprd = [] 
         B_S_dict = hashdict()#odict()
         if isinstance(z.port, InputPort) and isinstance(md, FSDEVS):
            drv_dx = md.rv_delta_x(z) 
            B_S_dict[hn] = drv_dx.table #deepcopy(drv_dx.table)
         elif isinstance(md, CoupledDEVS) or isinstance(z.port, OutputPort):
            XP = set()
            if isinstance(z.port, OutputPort): 
               md=md.parent
            if md:
               md.getAllInfluencee(z, XP)
               for xp in XP:
                  if isinstance(xp.parent, FSDEVS): xi = PortValue(xp, z.value)
                  drv_dx = xp.parent.rv_delta_x(xi)
                  B_S_dict[self.hiername(xp.parent)] = self.getDRV(drv_dx).table #deepcopy(getDRV(drv_dx).table)
            #print 'z=',z, 'XP=', XP
            #print 'B_S_dict=', B_S_dict

         crossprd = cross_product_dict_dict(B_S_dict, hashdict())#odict()) 
      return crossprd

   def get_nds_CR_p(self, influencee, nv, CR, p):
      '''get nds CR'''
      for k in influencee:
         (rho, ns) = influencee[k][0] # rho: schedule indicator, ns: next states
         #nv.set_ss(k, (ns, self.get_ta(ns,k))) #(deepcopy(ns), self.get_ta(ns,k))
         nv.s[k]=(ns, self.get_ta(ns,k)) #(deepcopy(ns), self.get_ta(ns,k))
         p *= influencee[k][1]
         if rho: 
            CR.add(k)

      return p
 
   def use_vtx_repo(self, vtx):
      '''use vertex repository'''
      vtx.use_repo()

   def use_edge_repo(self, agg_e):
      '''use edge repository'''
      if agg_e[Reachability.sr] not in self.edge_CR_repo:
         self.edge_CR_repo[agg_e[Reachability.sr]]=agg_e[Reachability.sr]
      agg_e[Reachability.sr] = self.edge_CR_repo[agg_e[Reachability.sr]]
      if agg_e[self.sTrg] not in self.edge_trg_z_repo:
         self.edge_trg_z_repo[agg_e[Reachability.sTrg]] = agg_e[Reachability.sTrg]
      agg_e[Reachability.sTrg] = self.edge_trg_z_repo[agg_e[Reachability.sTrg]]
      if agg_e[self.sp] not in self.edge_p_repo:
         self.edge_p_repo[agg_e[self.sp]] = agg_e[self.sp]
      agg_e[Reachability.sp] = self.edge_p_repo[agg_e[Reachability.sp]]

   def register_transition(self, G, srcv, nv, md, z, influencee, CR, p, V_T, Visited):
      '''register transition'''
      #print 'srcv=',srcv
      #md = z.port.parent
      hn = self.hiername(md)
      agg_e = hashdict() #{}        # aggregating edges
      agg_e[Reachability.sTrg]=(hn,str(z)) #e firing edge 

      ResetC=set()  # a set of influenced edges 
      for c in nv.s:
         if c in CR: ResetC.add(c)
         elif isInfinity(nv.s[c][1]): ResetC.add(c)

      if influencee == None or len(influencee)==0:
         agg_e[Reachability.sp]=p
      else:
         agg_e[Reachability.sp]=self.get_nds_CR_p(influencee, nv, ResetC, p)

      if isSameStateVectors(srcv.s, nv.s) and len(CR) == 0: # ignoring the event 
         #print 'ignoring the event'
         return
      
      agg_e[Reachability.sr]=tuple(ResetC)

      self.use_edge_repo(agg_e)

      schedules = get_schedules(nv.s)
      #print "state_vector=%s"%state_vector
      #print "schedules=%s"%schedules
      tmptid = TID().copy(nv.tz); 
      #print 'tmptid=%s'%tmptid
      #print 'schedules='
      #for c in schedules:
      #   print 'c=%s, schedules[c]=%s'%(c,schedules[c])

      nv.tz = tmptid.successor(ResetC, schedules)
      self.use_vtx_repo(nv)

      if nv.tz.valid() == False:
         print 'nv.tz=%s is invalid'%nv.tz
         raise Exception("Invalid Successor!")

      #print 'nv.tz=',str(nv.tz)
      #print 'nv=', nv

      #printV('%s, |Visited|='%line(), Visited)
      if nv not in Visited: 
         V_T.append(nv) # s needs to be tested
         #printV('%s, |V_T|='%line(), V_T)
         #print '#####################'
         #raw_input("return")
         if nv not in G.V: 
            G.addVtx(nv)
            if len(G.V)% 100 == 0: print ".",
            if len(G.V)%1000 == 0: 
               print '\n%d'%len(G.V)
            if len(G.V)%10000 == 0:
               print str(datetime.now())

      G.addEdge(srcv, nv, agg_e )
      #print "state_vector=%s, SN=%s"%(state_vector, nv.s)

   def expand_external_transitions(self, G, v, z, CR, p, V_T, Visited):
      '''get all z edges from the state s'''
      inflencee_vectors = self.get_state_vector_of_influencee(v.s, z)
      #printV('%s, |influencee_vectors|='%line(), inflencee_vectors)
      for influencee in inflencee_vectors: 
         self.register_transition(G, v, ReachableVtx(v.s, v.tz), z.port.parent, z, influencee, CR, p, V_T, Visited)

      #printV('%s, |influencee_vectors|='%line(), inflencee_vectors)
      #print "################ IS IT THE SAME ??????????????????????? "
      #raw_input("return")

   def expand_internal_transitions(self, G, v, V_T, Visited):
      '''expand enabled y edges from the state v'''
      for hn in v.s:
         #print 'v=%s',v
         if self.IsEnabledY(v.tz, v.s[hn][1], hn):
            #raw_input("%s is enabled"%hn)
            self.set_atomic_model_state(v.s)  
            md = self.AM[hn]
            drv_dy = md.rv_delta_y()
            #print 'drv_dy=', drv_dy
            table = self.getDRV(drv_dy).table
            for (y,ns) in table: # for each (y,ns)
               #md.set_state(v.s[hn]['s']) # reset the state of md
               p = table[(y,ns)]
               ymdhn = hn #hiername(y.port.parent)
               CR = set(); CR.add(ymdhn)
               inflencee_vectors = self.get_state_vector_of_influencee(v.s, y)
               if len(inflencee_vectors)>0:
                  for influencee in inflencee_vectors: 
                     nv = ReachableVtx(v.s, v.tz)
                     nv.tz.Enabler(ymdhn, v.s[ymdhn][1])
                     #print "nv.tz=%s"%nv.tz
                     if nv.tz.valid()==False:
                        raise Exception("nv.cv. is invalid")

                     #nv.set_ss(ymdhn,(ns, self.get_ta(ns, ymdhn)))
                     nv.s[ymdhn]=(ns, self.get_ta(ns, ymdhn))
                     self.register_transition(G, v, nv, md, y, influencee, CR, p, V_T, Visited) 
               else:
                  nv = ReachableVtx(v.s, v.tz)
                  nv.tz.Enabler(ymdhn, v.s[ymdhn][1])
                  #nv.set_ss(ymdhn, (ns, self.get_ta(ns, ymdhn)))
                  nv.s[ymdhn]=(ns, self.get_ta(ns, ymdhn))
                  self.register_transition(G, v, nv, md, y, {}, CR, p, V_T, Visited) 
         else:
            pass
            #raw_input("%s is not enabled"%hn)

   def getAllAtomicModels(self, md):
      '''get all FSDEVS models into self.AM'''
      if isinstance(md, FSDEVS):
         self.AM[self.hiername(md)] = md
      elif isinstance(md, CoupledDEVS):
         for kid in md.M:
            self.getAllAtomicModels(kid)

   def computes0Probability(self, state_vector):
      '''compute s0 probability'''
      tp = 1 
      for hn in state_vector:
         s0 = state_vector[hn][0]#['s']
         rv_s0 = self.AM[hn].rv_delta_0()
         table = self.getDRV(rv_s0).table
         if s0 in table:
            tp *= table[s0]
         else:
            tp *= 0
      return tp

   def GenReachFromFSD(self, md, X):
      '''generate all reachable state graph of and atomic FSD md'''
      G = Graph(); Visited = set(); V_T=[]

      # step 0.
      self.AM={}
      self.getAllAtomicModels(md)

      # step 1. 
      for discs0 in self.Delta0(): 
         #print "s0=", discs0,", p=",self.computes0Probability(discs0)
         schedules = get_schedules(discs0)
         tmptid = TID(); tmptid.makeZero(schedules)
         tid = tmptid.successor(set(), schedules)
         #print 'tid=', tid
         v = ReachableVtx(discs0, tid)
         G.addVtx(v, {'*':1} ); 
         #print 'v0=', v
         V_T.append(v)

      #print "|G.V0|=", len(G.V)
      #print G.vtxs2str()

      # step 2. for each state s in V_T, extend reachable states of s.
      while len(V_T)>0:  
         # step 2.1. Pick one state s from V_T and mark as visited 
         v = V_T.pop(); 
         Visited.add(v) # pick one state s, and mark s as visited

         # step 2.2. Add all internal transition edges from s
         # prepare the state of md as s 

         self.expand_internal_transitions(G, v, V_T, Visited)

         #print 'enabledE=', enabledE
         #print 'V_T=', V_T

         #-- step 2.3. Add all external transition edges from s and all x in X
         for x in X:
            self.expand_external_transitions(G, v, x, set(), 1, V_T, Visited)
            #print 'V_T=', V_T

      return G 
   
   def __str__(self):
      Visited = dict()
      V_T = [] 
   
      #print '|V|=', len(self.G.V)
      for v in self.G.V:
         if '*' in self.G.V[v]: # add the initial vertex v into V_T
            V_T.append(v)

      visited_order = 1

      strn = "\n############### Edges ###############\n"
      nE = 0
      while (len(V_T)>0):
         src = V_T.pop() 
         if src not in Visited:
            Visited[src] = visited_order; visited_order += 1
         
         if src in self.G.FE:
            for de in self.G.FE[src]: 
               nE += 1
               if S_dst in de: dest = de[S_dst]
               edge = dict(de); del edge[S_dst]; del edge[S_src]

               if dest not in Visited:
                  Visited[dest] = visited_order; visited_order += 1; V_T.append(dest)

               if len(edge)>0 :
                  strn += str(Visited[src])+" -[ "+dict_to_s(edge)+" ]-> "+str(Visited[dest])+"\n"
               else:
                  strn += str(Visited[src])+"--> "+str(Visited[dest])+"\n"
      strn += "\n####### Summary: |V|=%d, |E|=%d #######\n"%(len(Visited),nE)


      strnv = "\n########### Vertices #############\n"#%len(self.G.V)
      sorted_by_value = sorted(Visited.items(), key=lambda x: x[1])
      for v in sorted_by_value:
         if '*' in self.G.V[v[0]]:
            strnv += str(v[1])+ "*,\t" + str(v[0]) + "\n"
         else:
            strnv += str(v[1])+ ",\t" + str(v[0]) + "\n"

      strn = strnv + strn
      return strn


   ############################################################################################

if __name__ == '__main__':
   try:
      G = Reachability(None)
      print "Done!G=%"%str(G)
   except:
      print "Exception" 

